#include "non_max_suppression_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

#include <vector>
#include <cstring>
#include <cstdlib>
#include <algorithm>
namespace optiling {
const uint32_t BLOCK_SIZE = 32;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    NonMaxSuppressionTilingData tiling;
    int32_t NUM = 24;
    uint32_t sizeofdatatype;
    uint32_t totalLengthAligned;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);   //获取平台整个UB的UB的大小
 
    auto aivNum = ascendcPlatform.GetCoreNum();     //获取平台的总的核的数量
    auto scores_shape = context->GetInputTensor(1)->GetOriginShape();
    auto dt = context->GetInputTensor(0)->GetDataType();
    if(dt == ge::DT_UINT8){
        sizeofdatatype = 1;  //int8,uint8占1个字节
        // NUM =12;
    }else if(dt == ge::DT_FLOAT16 || dt == ge::DT_BF16){
        sizeofdatatype = 2;  //float16/half,int16,占2个字节
        // NUM=6;
    }else{
        sizeofdatatype = 4;  //float,int32占4个字节
        // NUM=6;
    }

    int32_t num_batches = scores_shape.GetDim(0);
    int32_t num_classes = scores_shape.GetDim(1);
    int32_t spatial_dimension = scores_shape.GetDim(2);
    std::cout<<"spatial_dimension:"<<spatial_dimension<<std::endl;
    auto max_output = context->GetInputTensor(2)->GetOriginShape();
    auto iouth_shape = context->GetInputTensor(3)->GetOriginShape();
    auto scoreth_shape = context->GetInputTensor(4)->GetOriginShape();

    uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();

    uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;   //每个数据类型所占的元素个数,即一个datablock里面存储的数据数
    uint32_t tiling_size = ((ub_size) / BLOCK_SIZE / 2) / NUM;    //平台总的ub大小/32字节/doublebuffer/队列数
    tiling_size = tiling_size <= 8 ? tiling_size : tiling_size / 8 * 8;  //计算的总的数据块数

    uint32_t block_size = tiling_size * ALIGN_NUM;  //计算单元处理的元素个数：half:128个元素，float:64个元素
    aivNum = (aivNum < totalLength / block_size) ? aivNum : (totalLength / block_size); //half:16个核，float：32个核,case1,256/64 = 4核
    aivNum = aivNum >= 1 ? aivNum : 1;

    uint32_t core_size = (totalLength / aivNum) / (ALIGN_NUM * 8) * (ALIGN_NUM * 8);  //分到计算核中的元素个数，half:128，float：64
    uint32_t core_remain = totalLength - aivNum * core_size;    //分好完整的核之后，剩余的元素数量
    int center_point_box = *context->GetAttrs()->GetInt(0);//0或1
    std::cout<<"center_point_box:"<<center_point_box<<std::endl;

    tiling.set_totalLength(totalLength);
    tiling.set_ALIGN_NUM(ALIGN_NUM);
    tiling.set_tiling_size(tiling_size);
    tiling.set_block_size(block_size);
    tiling.set_aivNum(aivNum);
    tiling.set_core_size(core_size);
    tiling.set_core_remain(core_remain);
    tiling.set_center_point_box(center_point_box);
    tiling.set_num_batches(num_batches);
    tiling.set_num_classes(num_classes);
    tiling.set_spatial_dimension(spatial_dimension);
    context->SetBlockDim(aivNum);

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class NonMaxSuppression : public OpDef {
public:
    explicit NonMaxSuppression(const char* name) : OpDef(name)
    {
        this->Input("boxes")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("scores")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("max_output_boxes_per_class")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("iou_threshold")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("score_threshold")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("selected_indices")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Attr("center_point_box").AttrType(OPTIONAL).Int(0);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(NonMaxSuppression);
}
